Add bounded-memory training and colocation controls#1890
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces several memory optimization features, including CPU-resident microbatches, sequence-dimension chunking for Megatron vocab entropy computation, and a colocated worker memory barrier with optional hard-eviction for inactive reference workers. The review feedback suggests mapping PyTorch device indices to physical GPU indices via CUDA_VISIBLE_DEVICES to ensure accurate NVML memory statistics, and removing a fragile .squeeze(0) call on entropy chunks to robustly support varying input dimensions.
9b6a19a to
8af9294
Compare
8af9294 to
782a8c7
Compare
erictang000
left a comment
erictang000
left a comment
There was a problem hiding this comment.
looks good, just some minor comments
| return torch.cat(entropy_chunks, dim=-1) | ||
|
|
||
|
|
||
| def vocab_parallel_entropy_weighted_sum( |
There was a problem hiding this comment.
just to confirm these changes are stable, can you run a basic gsm8k example before and after this PR, and link it here? Just want to check there is no significant speed regression from chunking entropy calculation, and that the entropy metric is exactly matching before
There was a problem hiding this comment.
We have a unit test on this part: tests/backends/skyrl_train/distributed/test_vocab_entropy_chunking.py, this verifies that the chunked CE is numerically identical. For the overhead, the following table shows:
chunk_memory_mb |
tokens/chunk | #chunks | exact | op time | vs unchunked |
|---|---|---|---|---|---|
| off (unchunked) | — | — | — | 6.3 ms | 1.0× |
| 4096 | 4096 | 2 | ✓ | 6.4 ms | 1.0× |
| 2048 | 2048 | 4 | ✓ | 7.5 ms | 1.2× |
| 1024 | 1024 | 8 | ✓ | 14.9 ms | 2.4× |
| 512 (current default) | 512 | 16 | ✓ | 29.7 ms | 4.7× |
| 256 | 256 | 32 | ✓ | 59.6 ms | 9.4× |
| 128 | 128 | 64 | ✓ | 118.6 ms | 18.8× |
The slowdown is expected, but since the overall step time here is at most hundreds of ms, the impact of chunking is almost negligible.
PPORayActorGroup stored `pg` and `num_gpus_per_actor` on self but nothing ever read them back (the constructor already passes both directly to _initiate_actors). Drop the two dead assignments; the constructor params are unchanged.
Make the previously opt-in `cpu_resident_microbatches` behavior the only path and remove the config flag. Each microbatch is now moved to the device just before its forward step (in MegatronModelWrapper.forward_step), instead of moving the whole mini-batch up front. This caps resident input-tensor memory at one microbatch and is a no-op when a tensor is already on device. The per-microbatch transfer is the single choke point for every forward and forward-backward pass, so the up-front `.to(device)` calls in the worker are removed; nothing downstream reads the mini-batch tensors on-device after the microbatch loop (only chunk()/BatchIterator, which slice on CPU, and metadata). Removes the MegatronConfig field, the wrapper constructor param, the worker plumbing, and the CLI-override test for the flag.
1d15433 to
547d120
Compare
erictang000
left a comment
erictang000
left a comment
There was a problem hiding this comment.
nice, looking good thanks!
3 participants
Bounded-Memory Megatron Training and Colocation
Summary
This PR adds controls that bound temporary training memory and make colocated Megatron offload behavior complete.
Changes
CPU-resident policy microbatches is opt-in; automatic vocabulary-entropy chunk sizing, which defaults to a 512 MiB temporary-memory budget.
Memory Validation